#include "ESPUpdateServer.h"
#include "StreamString.h"
#include "compat.h"

static const char serverIndex[] PROGMEM =
    R"(<html><body><form method='POST' action='' enctype='multipart/form-data'>
                  <input type='file' name='update'>
                  <input type='submit' value='Update'>
               </form>
         </body></html>)";
static const char successResponse[] PROGMEM =
    "<META http-equiv=\"refresh\" content=\"15;URL=/\">Update Success! "
    "Rebooting...\n";

ESPUpdateServer::ESPUpdateServer(bool serial_debug) {
  _serial_output = serial_debug;
  _server = NULL;
  _username = emptyString;
  _password = emptyString;
  _authenticated = false;
}

void ESPUpdateServer::setup(ESP_SERVER_CLASS* server,
                            const String& path,
                            const String& username,
                            const String& password) {
  _server = server;
  _username = username;
  _password = password;

  // handler for the /update form page
  _server->on(path.c_str(), HTTP_GET, [&]() {
    if (_username != emptyString && _password != emptyString &&
        !_server->authenticate(_username.c_str(), _password.c_str()))
      return _server->requestAuthentication();
    _server->sendHeader("Connection", "close");
    _server->send_P(200, PSTR("text/html"), serverIndex);
  });

  // handler for the /update form POST (once file upload finishes)
  _server->on(
      path.c_str(), HTTP_POST,
      [&]() {
        if (!_authenticated)
          return _server->requestAuthentication();
        if (Update.hasError()) {
          _server->send(200, F("text/html"),
                        String(F("Update error: ")) + _updaterError);
        } else {
          _server->client().setNoDelay(true);
          _server->sendHeader("Connection", "close");
          _server->send_P(200, PSTR("text/html"), successResponse);
          delay(100);
          _server->client().stop();
          ESP.restart();
        }
      },
      [&]() {
        // handler for the file upload, get's the sketch bytes, and writes
        // them through the Update object
        HTTPUpload& upload = _server->upload();

        if (upload.status == UPLOAD_FILE_START) {
          _updaterError = String();
          if (_serial_output)
            Serial.setDebugOutput(true);

          _authenticated =
              (_username == emptyString || _password == emptyString ||
               _server->authenticate(_username.c_str(), _password.c_str()));
          if (!_authenticated) {
            if (_serial_output)
              Serial.printf("[OTA] Unauthenticated Update\n");
            return;
          }

          int command =
              (upload.filename.indexOf("spiffs") > -1) ? U_SPIFFS : U_FLASH;
          if (_serial_output)
            Serial.printf("[OTA] Update Start: %s\n", upload.filename.c_str());
          Serial.printf("[OTA] Update Type: %s\n",
                        command == U_FLASH ? "U_FLASH" : "U_SPIFFS");

#if defined(ESP8266)
          uint32_t maxSketchSpace =
              (ESP.getFreeSketchSpace() - 0x1000) & 0xFFFFF000;
          if (!Update.begin(maxSketchSpace)) {  // start with max available size
            _setUpdaterError();
          }
#elif defined(ESP32)
          if (!Update.begin(UPDATE_SIZE_UNKNOWN, command)) {
            // start with max available size
            Update.printError(Serial);
          }
#endif
        } else if (_authenticated && upload.status == UPLOAD_FILE_WRITE &&
                   !_updaterError.length()) {
          if (_serial_output)
            Serial.print(".");
          if (Update.write(upload.buf, upload.currentSize) !=
              upload.currentSize) {
            _setUpdaterError();
          }
        } else if (_authenticated && upload.status == UPLOAD_FILE_END &&
                   !_updaterError.length()) {
          // true to set the size to the current progress
          if (Update.end(true)) {
            if (_serial_output)
              Serial.printf("\n[OTA] Update Success: %u\nRebooting...\n",
                            upload.totalSize);
          } else {
            _setUpdaterError();
          }
          if (_serial_output)
            Serial.setDebugOutput(false);
        } else if (_authenticated && upload.status == UPLOAD_FILE_ABORTED) {
          Update.end();
          if (_serial_output)
            Serial.println("[OTA] Update was aborted");
        } else {
          Serial.printf("[OTA] Update Failed Unexpectedly: status=%d\n",
                        upload.status);
        }
        delay(0);
      });
}

void ESPUpdateServer::_setUpdaterError() {
  if (_serial_output)
    Update.printError(Serial);
  StreamString str;
  Update.printError(str);
  _updaterError = str.c_str();
}
